#include "AMG.h"
#include <algorithm>

using u8 = unsigned char;
using std::fabs;
using std::find;

void DealWithBoundary2(SpMat &Spt, VectorXd &F,const vector<pair<u32, double>> &Boundval)
{
    for (auto &p : Boundval)
    {
        u32 idx = p.first;
        for (SpMat::InnerIterator it(Spt, idx); it; ++it)
        {
            u32 row = it.row();
            if (row == idx)
                continue;
            F[row] -= it.value() * p.second;
            Spt.coeffRef(idx, row) = 0;
            Spt.coeffRef(row, idx) = 0;
        }
        F[idx] = Spt.coeffRef(idx, idx) * p.second;
    }
};
double L2_Error(VectorXd v1, VectorXd v2)
{
    double error1 = 0;
    double error2 = 0;
    for(int i = 0; i < v1.size(); i++)
    {
        error1 += v1[i] * v1[i] + v2[i] * v2[i];
        error2 += v1[i] * v2[i];
    }
    return sqrt(error1 - 2 * error2) / v1.size();
};
double L2ero_simp(VectorXd v1, VectorXd v2)
{
    double error = 0;
    for(int i = 0; i < v1.size(); i++)
    {
        error += (v1[i] - v2[i]) * (v1[i] - v2[i]);
    }
    return sqrt(error) / v1.size();
}
void Interator(SpMat &_A, VectorXd &_u, VectorXd &_Rhs, InteratorMethon _type, double _weight)
{
    if (_type == WJ)
    {
        VectorXd v = Eigen::VectorXd::Zero(_Rhs.size());
        for( int i =0; i < _u.size(); ++i)
        {
            double temp = _Rhs[i];
            double a_ii = _A.coeffRef(i,i);
            {
                for(Eigen::SparseMatrix<double>::InnerIterator it(_A,i); it; ++it)
                {
                    int i_row = it.row();
                    if(i_row == i)
                        continue;
                    else
                        temp -=(it.value() * (_u[i_row]));
                }
                v[i] = temp / a_ii;
            }
        }
        v = (1 - _weight) * _u + _weight * v ;
        _u = v ;
    }
    else if (_type == GS)
    {
        for( int i =0; i < _u.size(); ++i)
        {
            double temp = _Rhs[i];
            double a_ii = _A.coeffRef(i,i);
            {
                for(Eigen::SparseMatrix<double>::InnerIterator it(_A,i); it; ++it)
                {
                    int i_row = it.row();
                    if (i_row == i)
                        continue;
                    else
                        temp -= (it.value() * (_u[i_row]));
                }
                _u[i] = temp / a_ii;
            }
        }
    }
    else if (_type == SOR)
    {
        for(int i =0; i < _u.size(); ++i)
        {
            double temp = _Rhs[i];
            double a_ii = _A.coeffRef(i,i);
            {
                for(Eigen::SparseMatrix<double>::InnerIterator it(_A,i); it; ++it)
                {
                    int i_row = it.row();
                    if (i_row == i)
                        continue;
                    else
                        temp -= (it.value() * (_u[i_row]));
                }
                _u[i] = temp / a_ii * _weight + _u[i] * (1 - _weight);
            }
        }
    }
    else if (_type == SSOR)
    {
        for(int i =0; i < _u.size(); ++i)
        {
            double temp = _Rhs[i];
            double a_ii = _A.coeffRef(i,i);
            {
                for(Eigen::SparseMatrix<double>::InnerIterator it(_A,i); it; ++it)
                {
                    int i_row = it.row();
                    if (i_row == i)
                        continue;
                    else
                        temp -= (it.value() * (_u[i_row]));
                }
                _u[i] = temp / a_ii * _weight + _u[i] * (1 - _weight);
            }
        }
        for( int i = _u.size()-1; i >= 0; --i)
        {
            double temp = _Rhs[i];
            double a_ii = _A.coeffRef(i,i);
            {
                for(Eigen::SparseMatrix<double>::InnerIterator it(_A,i); it; ++it)
                {
                    int i_row = it.row();
                    if(i_row == i)
                        continue;
                    else
                        temp -= (it.value() * (_u[i_row]));
                }
                _u[i] = temp / a_ii * _weight + _u[i] * (1 - _weight);
            }
        }
    }
};

void AMG_Multigrid::V_Cycle()
{
    VectorXd u_init = VectorXd::Zero(F_Rhs[0].size());
    iter_num = 0 ;
    u32 VCycleNum = 0 ;
    u32 itermaxnum = 1000;
    u_Lhs[0] = u_init;
    double errorinit = L2_Error(F_Rhs[0], Vec_Mat[0]*u_Lhs[0]);
    cout << "init error is : " << errorinit << endl;
    errorinit *= Tolerance;
    auto t_start = std::chrono::high_resolution_clock::now();
    while(!Has_Solved && VCycleNum < itermaxnum){
        VCycleNum++;
        V_CycleOnce(0);
        double l2error = L2_Error(F_Rhs[0], Vec_Mat[0] * u_Lhs[0]);
        // cout << VCycleNum <<" : " << l2error  << endl;
        if(l2error < errorinit)
            Has_Solved = true;
    }
    if (Has_Solved)
        cout <<"Vcycle num is : " <<  VCycleNum << endl 
            <<"Relative error is: " << L2_Error(F_Rhs[0], Vec_Mat[0] * u_Lhs[0]) * Tolerance / errorinit << endl;
    else{
        cout << "after " <<itermaxnum <<" times iterator, relative error is : " << endl
            << L2_Error(F_Rhs[0], Vec_Mat[0] * u_Lhs[0]) * Tolerance / errorinit << endl;
    }
    auto t_end = std::chrono::high_resolution_clock::now();
    float t_total = std::chrono::duration<float, std::milli>(t_end - t_start).count();
    // cout <<"calculate over!\n" << "The number of iterator is :" << iter_num << endl;
    cout << "Solving Equation takes " << t_total <<" ms" << endl;
    // cout <<"V_cycle number is :" << VCycleNum << endl;
};

void AMG_Multigrid::V_CycleOnce(int _NowLayer)
{
    if(_NowLayer == N_Layer -1)        
    {
        
        VectorXd &f = F_Rhs[_NowLayer];
        
        // 最粗网格，采用直接法求解
        u_Lhs[_NowLayer] = LastmatLU.solve(f);
	
        return ;
    }
    else{
        int vis1 = 4;
        VectorXd &u = u_Lhs[_NowLayer];
        SpMat &A = Vec_Mat[_NowLayer];
        VectorXd &f = F_Rhs[_NowLayer];
        //先在当前网格迭代vis1次
        for(int i =0; i < vis1; i++)
        {
            Interator(A, u, f, InteratorType, InteratorWeight);
            // cout << L2_Error(A * u, f);
            iter_num++ ;
        }
        VectorXd r_h = f - A *u ;
        VectorXd r_2h = Restrict[_NowLayer]* r_h;         
        F_Rhs[_NowLayer + 1] = r_2h ;
        
        //初始迭代值为零向量
        VectorXd zero_u_2h = VectorXd::Zero(r_2h.size());
        u_Lhs[_NowLayer + 1] = zero_u_2h;
        V_CycleOnce(_NowLayer + 1);
        VectorXd e_h = Interpolation[_NowLayer]* u_Lhs[_NowLayer + 1];     
        u = u + e_h ;
        for(int i =0; i < vis1; i++)
        {
            Interator(A, u, f, InteratorType, InteratorWeight) ;
            iter_num++ ;
        }
        u_Lhs[_NowLayer] = u;
    }
};

void AMG_Multigrid::Set_InteratorType(InteratorMethon _InteratorType, double _Weight)
{   
    InteratorType = _InteratorType;
    InteratorWeight = _Weight;
};

void AMG_Multigrid::Set_Tolerance(double _Tolerance)
{
    Tolerance = _Tolerance;
};
/// @brief 
/// @param _Spt 求解的矩阵
/// @param _Rhs 求解的右端项
/// @param _Tolerance 迭代终止的阈值
/// @param _InteratorType 选用的迭代法：WJ, GS, SOR, SSOR
/// @param _InteratorWeight 选用的WJ迭代法的权重
/// @param _theta 选取强连接点的阈值
AMG_Multigrid::AMG_Multigrid(SpMat &_Spt, VectorXd &_Rhs, double _Tolerance ,InteratorMethon _InteratorType ,double _InteratorWeight, double _theta)
{
    double EndRatio = 0.8;
    Tolerance = _Tolerance;
    InteratorType = _InteratorType;
    InteratorWeight = _InteratorWeight;
    Theta = _theta;
    N_Layer = 1;
    Vec_Mat.reserve(10);
    F_Rhs.reserve(10);
    Restrict.reserve(10);
    Interpolation.reserve(10);
    auto t_start = std::chrono::high_resolution_clock::now();
    Vec_Mat.push_back(_Spt);
    F_Rhs.push_back(_Rhs);
    while (1.0 * Vec_Mat.back().nonZeros() / Vec_Mat.back().size() < EndRatio)// || Vec_Mat.back().size() > 250000)
    {
        //生成插值算子
        BuildAMG bdamg(Vec_Mat.back(), Theta);
        bdamg.Build();
        SpMat itplt = bdamg.GetItplatMat();
        //保存插值算子以及投影算子
        Interpolation.push_back(itplt);
        Restrict.push_back(itplt.transpose());
        SpMat Sp2 = Restrict.back() * Vec_Mat.back() * itplt;  
        Vec_Mat.push_back(Sp2);
        F_Rhs.push_back(VectorXd::Zero(Vec_Mat.back().rows()));
        N_Layer++;
    }
    
    Eigen::MatrixXd mat(Vec_Mat.back());
    
    LastmatLU = Eigen::PartialPivLU<Eigen::MatrixXd> (mat);
    cout <<"number of dofs : " << Vec_Mat[0].rows() << endl;
    
    u_Lhs.resize(N_Layer);
    auto t_end = std::chrono::high_resolution_clock::now();
    float t_total = std::chrono::duration<float, std::milli>(t_end - t_start).count();
    cout << "Build matrix over! It takes " << t_total <<" ms" << endl;
}

void BuildAMG::PreBuild()
{
    ///////////首先，获得Sm,St,Sw 分别保存第i个点强连接，与第i个点强连接，与第i个点非强连接的点
    u32 n = A.cols();
    vector<vector<u32>> Sm(n);
    vector<vector<u32>> St(n);
    beta.resize(n);
    ////////////////////注意，eigen中的矩阵是按列压缩的，但是因为所有的矩阵都是对称阵，所以按列读取也可以
    for (u32 i = 0; i < n; i++)
    {
        beta[i] = 0;
        for (SpMat::InnerIterator it(A, i); it; ++it)
        {
            u32 row = it.row();
            if (row == i)
                continue;
            if (beta[i] < fabs(it.value()))
                beta[i] = fabs(it.value());
        }
        beta[i] *= theta;
        for (SpMat::InnerIterator it(A, i); it; ++it)
        {
            u32 row = it.row();
            if (row == i)
                continue;
            if (fabs(it.value()) > beta[i])
            {
                Sm[i].push_back(row);
                St[row].push_back(i);
            }
        }
    }
    vector<int> lambdaI(n);
    for (u32 i = 0; i < n; i++)
        lambdaI[i] = St[i].size();
    SegmentTree tree(lambdaI);
    vector<u8> Um(n, '0');          //0代表未确定的点，1代表粗网格，2代表细网格
    while (tree.getmax().second != 0)
    {
        u32 idx = tree.getmax().first;
        Um[idx] = '1';
        tree.update(idx, 0);
        for (auto &j : St[idx])
        {
            if (Um[j] == '0')
            {
                Um[j] = '2';
                tree.update(j, 0);
                for (auto &l : Sm[j])
                {
                    if (Um[l] == '0')
                        tree.add(l, 1);
                }
            }
        }
        for (auto &j : Sm[idx])
            if (Um[j] == '0')
                tree.add(j, -1);
    }
    Cpoints.clear();
    for (u32 i = 0; i < n; i++)
    {
        if (Um[i] == '1'){
            Cpoints.push_back(i);
        }
    }
}
void BuildAMG::buildall()
{
    
    u32 n = A.rows();
    vector<int> f2c(n, -1);
    for (u32 i = 0; i < Cpoints.size(); i++)
        f2c[Cpoints[i]] = i;
    ItplatMat = SpMat(n, Cpoints.size());
    vector<Tri> Trilist;

    for (u32 i = 0; i < n; i++)
    {
        ///////////对于一个点，若它同时是粗网格和细网格的点，那么这个点上的误差插值就是它本身///////////
        if (f2c[i] != -1)
        {
            // ItplatMat.insert(i, f2c[i]) = 1;
            Trilist.push_back(Tri(i, f2c[i], 1));
            continue;
        }
        double SumAin = A.coeffRef(i, i);
        vector<u32> Fs;     //i周围的细网格的强连接点
        vector<u32> CI;
        
        for (SpMat::InnerIterator it(A, i); it; ++it)
        {
            u32 row = it.row();
            if (row == i)
                continue;
            if (f2c[row] != -1)
                CI.push_back(row);
            else{
                if (fabs(it.value()) < beta[i])
                {
                    SumAin += it.value();
                }    
                else
                    Fs.push_back(row);
            }
        }
        
        vector<double> SumAmk(Fs.size(), 0);
        for (u32 j = 0; j < Fs.size(); ++j)
        {
            u32 m = Fs[j];
            for (SpMat::InnerIterator it(A, m); it; ++it)
            {
                u32 row = it.row();
                if (find(CI.begin(), CI.end(), row) != CI.end())
                    SumAmk[j] += it.value();
            }
        }

        for (u32 j = 0; j < Fs.size(); ++j)
        {
            u32 m = Fs[j];
            for (SpMat::InnerIterator it(A, m); it; ++it)
            {
                u32 row = it.row();
                if (find(CI.begin(), CI.end(), row) != CI.end())
                    if (std::fabs(SumAmk[j]) > 1e-12)
                        Trilist.push_back(Tri(i, f2c[row], -1 * A.coeffRef(i, m) * it.value() / (SumAmk[j] * SumAin)));
            }
        }
        for (auto &v : CI)
        {
            Trilist.push_back(Tri(i, f2c[v], -1 * A.coeffRef(i, v) / SumAin));
        }
        
    } 
    ItplatMat.setFromTriplets(Trilist.begin(), Trilist.end());
    ItplatMat.makeCompressed();
};

